package cn.xyfmzmy.flyweight;

import java.nio.ByteBuffer;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 模拟 Kafka 的内存池,机制为：
 * 1. 池子的总大小为 totalSize，每个槽的大小为 slotSize。
 * 2. 一个程序正好申请了slotSize大小的空间，那么归还后，这个空间会被标记可复用的 slot
 * 当其他程序申请slotSize大小的空间时，就会把可复用的 slot 直接返回给程序
 * 3. 假设有一个程序申请了一个较大的空间，内存池剩余的空间不够分配了，但是加上可复用的 slot 空间后就够了的话
 * 就会把可复用的 slot 删除，归还到剩余的空间中，然后分配给程序。
 */
public class BufferPool {

    private final Integer totalSize;        // 内存池总大小
    private final Integer slotSize;         // 每个槽的大小
    private long free;                      // 内存池的剩余空间
    private final Deque<ByteBuffer> slotQueue = new ArrayDeque<>();     // 存储可复用的 slot
    private final Deque<Condition> waiters = new ArrayDeque<>();        // 存储正在等待空间释放的任务
    private final Lock lock = new ReentrantLock();      // 可重入锁

    public BufferPool(Integer totalSize, Integer slotSize) {
        this.totalSize = totalSize;
        this.slotSize = slotSize;
        this.free = totalSize;
    }

    public ByteBuffer allocate(int size,long timeout) throws InterruptedException {
        if(size > totalSize || size < 0){
            throw new RuntimeException("申请空间大小异常");
        }
        lock.lock();
        try {
            // 申请的正好是一个槽的大小，并且可复用的 slot 不为空
            if(size == slotSize && !slotQueue.isEmpty()){
                return slotQueue.pollFirst();
            }
            // 剩余空间 + 可复用槽的空间 >= 申请的空间，先释放 slot 再分配
            if((free + slotSize * slotQueue.size()) >= size){
                freeUp(size);
                free -= size;
                return ByteBuffer.allocate(size);
            }
            // 空闲空间不够分配了，等待其它线程释放空间
            Condition condition = lock.newCondition();
            waiters.addLast(condition);
            long remainTime = timeout;
            try {
                while(true){
                    long start = System.currentTimeMillis();
                    // 如果是被唤醒的，wakeUp为true; 如果等待超时，wakeUp为false
                    boolean wakeUp = condition.await(remainTime, TimeUnit.MILLISECONDS);
                    if(!wakeUp){
                        throw new RuntimeException("分配失败，等待释放超时");
                    }
                    // 成功被唤醒，同样判断剩余的空间是否够分配
                    if(size == slotSize && !slotQueue.isEmpty()){
                        return slotQueue.pollFirst();
                    }
                    if((free + slotSize * slotQueue.size()) >= size){
                        freeUp(size);
                        free -= size;
                        return ByteBuffer.allocate(size);
                    }
                    // 还是不够分配,更新等待的剩余时间，重新等待空间被释放
                    remainTime -= System.currentTimeMillis() - start;
                }
            } finally {
                waiters.remove(condition);
            }
        } finally {
            lock.unlock();
        }
    }


    // 释放可复用的slot，直到剩余空间足够分配size大小的空间
    private void freeUp(int size){
        while(free < size && !slotQueue.isEmpty()){
            free += slotQueue.pollFirst().capacity();
        }
    }

    public void deallocate(ByteBuffer byteBuffer){
        lock.lock();
        try {
            // 释放的空间正好是一个槽的大小，标记为可复用的槽
            if(byteBuffer.capacity() == slotSize){
                slotQueue.addLast(byteBuffer);
            } else {
                // 其他情况，直接归还空间
                free += byteBuffer.capacity();
            }
            if(!waiters.isEmpty()){
                // 唤醒等待空间释放的线程
                waiters.peekFirst().signal();
            }
        } finally {
            lock.unlock();
        }
    }


}
